#include "cinn/ir/ir.h"

#include <map>
#include <string>
#include <vector>

#include "cinn/common/cinn_value.h"
#include "cinn/common/ir_util.h"
#include "cinn/ir/ir_printer.h"
#include "cinn/ir/ir_visitor.h"
#include "cinn/ir/tensor.h"
#include "cinn/lang/module.h"
#include "cinn/optim/ir_simplify.h"

namespace cinn {
namespace ir {

using common::make_shared;

Expr Cast::Make(Type t, Expr v) {
  CHECK(!t.is_unk());
  CHECK(!(t.is_void() && !t.is_cpp_handle())) << "Void is not allowed to cast";
  CHECK(v.defined());

  auto node = make_shared<Cast>();
  node->v() = v;
  node->set_type(t);
  return Expr(node);
}

void Cast::Accept(IRVisitor *v) const { v->IRVisitorBase::Visit(&this->v()); }

Expr Add::Make(Expr a, Expr b) {
  auto node = make_shared<Add>(a, b);
  return Expr(node);
}

Add::Add(Expr a, Expr b) : BinaryOpNode<Add>(a.type(), a, b) {}

Expr Sub::Make(Expr a, Expr b) {
  auto node = make_shared<Sub>(a, b);
  return Expr(node);
}

Expr Mul::Make(Expr a, Expr b) {
  CHECK(a.defined());
  CHECK(b.defined());
  CHECK_EQ(a.type(), b.type());
  auto node = make_shared<Mul>(a, b);
  return Expr(node);
}

Expr Div::Make(Expr a, Expr b) {
  auto node = make_shared<Div>(a, b);
  return Expr(node);
}

Expr Mod::Make(Expr a, Expr b) {
  auto node = make_shared<Mod>(a, b);
  return Expr(node);
}

Expr Min::Make(Expr a, Expr b) {
  auto node = make_shared<Min>(a, b);
  return Expr(node);
}

Expr Max::Make(Expr a, Expr b) {
  auto node = make_shared<Max>(a, b);
  return Expr(node);
}

Expr Minus::Make(Expr a) {
  auto node = make_shared<Minus>(a);
  return Expr(node);
}

Expr EQ::Make(Expr a, Expr b) {
  auto node = make_shared<EQ>(a, b);
  return Expr(node);
}

Expr NE::Make(Expr a, Expr b) {
  auto node = make_shared<NE>(a, b);
  return Expr(node);
}

Expr LT::Make(Expr a, Expr b) {
  auto node = make_shared<LT>(a, b);
  return Expr(node);
}

Expr LE::Make(Expr a, Expr b) {
  auto node = make_shared<LE>(a, b);
  return Expr(node);
}

Expr GT::Make(Expr a, Expr b) {
  auto node = make_shared<GT>(a, b);
  return Expr(node);
}

Expr GE::Make(Expr a, Expr b) {
  auto node = make_shared<GE>(a, b);
  return Expr(node);
}

Expr And::Make(Expr a, Expr b) {
  auto node = make_shared<And>(a, b);
  return Expr(node);
}

Expr Or::Make(Expr a, Expr b) {
  auto node = make_shared<Or>(a, b);
  return Expr(node);
}

Type Or::type() const { return type_; }

Expr Not::Make(Expr v) {
  auto node = make_shared<Not>(v);
  return Expr(node);
}

Type Not::type() const { return type_; }

Expr Let::Make(Expr symbol, Expr body) {
  auto *n = make_shared<Let>();
  CHECK(symbol.type().valid());
  if (body.defined()) {
    CHECK(body.type().valid());
  }
  n->symbol = symbol;
  n->body   = body;
  n->set_type(n->symbol->type());
  return Expr(n);
}

Type Let::type() const { return symbol.type(); }

Expr _Var_::Make(const std::string &name, const Type &type) {
  auto node = new _Var_(name, type);
  return Expr(node);
}

Expr _Var_::Make(Expr lower_bound, Expr upper_bound, const std::string &name) {
  auto *n           = make_shared<_Var_>();
  n->lower_bound    = lower_bound;
  n->upper_bound    = upper_bound;
  n->is_reduce_axis = true;
  n->name           = name;
  n->set_type(lower_bound.type());
  return Expr(n);
}

Expr _Var_::Copy() const {
  auto *n           = make_shared<_Var_>();
  n->name           = name;
  n->is_reduce_axis = is_reduce_axis;
  n->lower_bound    = lower_bound;
  n->upper_bound    = upper_bound;
  n->set_type(type());
  return Expr(n);
}

Expr For::Make(
    Var loop_var, Expr min, Expr extent, ForType for_type, DeviceAPI device_api, Expr body, VectorizeInfo vector_info) {
  auto node = make_shared<For>();
  CHECK(loop_var.defined());
  CHECK(min.defined());
  CHECK(extent.defined());
  node->loop_var   = loop_var;
  node->min        = min;
  node->extent     = extent;
  node->device_api = device_api;
  node->body       = body;
  node->set_for_type(for_type);
  node->set_vectorize_info(vector_info);

  if (node->is_vectorized()) CHECK(node->vectorize_info().valid());
  return Expr(node);
}

std::vector<Expr *> For::expr_fields() { return {&min, &extent, &body}; }
std::vector<const Expr *> For::expr_fields() const { return {&min, &extent, &body}; }

Expr Block::Make(const std::vector<Expr> &stmts) {
  auto node   = make_shared<Block>();
  node->stmts = stmts;
  return Expr(node);
}
std::vector<Expr *> Block::expr_fields() {
  std::vector<Expr *> res;
  for (auto &x : stmts) res.push_back(&x);
  return res;
}
std::vector<const Expr *> Block::expr_fields() const {
  std::vector<const Expr *> res;
  for (auto &x : stmts) res.push_back(&x);
  return res;
}

Expr IfThenElse::Make(Expr condition, Expr true_case, Expr false_case) {
  auto node = make_shared<IfThenElse>(condition, true_case, false_case);
  return Expr(node);
}

IfThenElse::IfThenElse(Expr condition, Expr true_case, Expr false_case)
    : ExprNode(Type()), condition(condition), true_case(true_case), false_case(false_case) {
  CHECK(condition.defined());
  CHECK(true_case.defined());
}
std::vector<Expr *> IfThenElse::expr_fields() { return {&condition, &true_case, &false_case}; }
std::vector<const Expr *> IfThenElse::expr_fields() const { return {&condition, &true_case, &false_case}; }

Expr Store::Make(Expr tensor, Expr value, const std::vector<Expr> &indices) {
  CHECK(tensor.As<_Tensor_>()) << "tensor should be _Tensor_ type";
  auto node     = make_shared<Store>();
  node->tensor  = tensor;
  node->value   = value;
  node->indices = indices;

  for (auto &indice : indices) {
    if (indice.As<Add>()) {
      if (indice.As<Add>()->b().As<Ramp>() || indice.As<Add>()->a().As<Ramp>()) {
        LOG(FATAL) << "found";
      }
    }
  }
  if (tensor->type() != Void()) {
    node->set_type(tensor->type().ElementOf().with_lanes(node->index().type().lanes()));
  }
  return Expr(node);
}

Expr Store::index() const {
  auto *tensor_n = tensor.As<ir::_Tensor_>();
  CHECK(tensor_n);
  Expr res = common::IndiceToAbsOffset(tensor_n->shape, indices);
  optim::Simplify(&res);
  return res;
}

Type Store::type() const { return value.type(); }
std::vector<Expr *> Store::expr_fields() {
  std::vector<Expr *> exprs({&tensor, &value});
  for (auto &idx : indices) exprs.push_back(&idx);
  return exprs;
}

std::vector<const Expr *> Store::expr_fields() const {
  std::vector<const Expr *> exprs({&tensor, &value});
  for (auto &idx : indices) exprs.push_back(&idx);
  return exprs;
}

Expr Alloc::Make(Expr dest, Type type, const std::vector<Expr> &extents, Expr condition, Expr body) {
  auto node = make_shared<Alloc>();
  CHECK(dest.As<_Buffer_>()) << "Alloc destination only supports Buffer";
  node->destination = dest;
  node->extents     = extents;
  node->condition   = condition;
  node->body        = body;
  node->set_type(type);
  return Expr(node);
}

int32_t Alloc::ConstantAllocationSize() const { return ConstantAllocationSize(extents); }

int32_t Alloc::ConstantAllocationSize(const std::vector<Expr> &extents) {
  int32_t res{1};
  for (auto &e : extents) {
    auto *p = e.As<IntImm>();
    CHECK(p) << "extent should be IntImm";
    res *= p->value;
  }
  return res;
}
std::vector<Expr *> Alloc::expr_fields() {
  std::vector<Expr *> res;
  for (auto &x : extents) res.push_back(&x);
  res.push_back(&condition);
  res.push_back(&body);
  return res;
}
std::vector<const Expr *> Alloc::expr_fields() const {
  std::vector<const Expr *> res;
  for (auto &x : extents) res.push_back(&x);
  res.push_back(&condition);
  res.push_back(&body);
  return res;
}

Expr Free::Make(Expr dest) {
  auto node = make_shared<Free>();
  CHECK(dest.As<_Buffer_>()) << "Free destination only supports Buffer";
  node->destination = dest;
  return Expr(node);
}

void _Range_::Accept(IRVisitor *v) const { v->Visit(this); }

Range::Range(_Range_ *n) : IrNodeRef(n) {}

Expr Call::Make(Type type,
                const std::string &name,
                const std::vector<Expr> &read_args,
                const std::vector<Expr> &write_args,
                CallType call_type,
                FunctionRef func,
                int value_index) {
  for (size_t i = 0; i < read_args.size(); ++i) {
    CHECK(read_args[i].defined());
  }

  auto node         = common::make_shared<Call>(type);
  node->name        = name;
  node->read_args   = read_args;
  node->write_args  = write_args;
  node->call_type   = call_type;
  node->func        = func;
  node->value_index = value_index;
  node->set_type(type);
  return Expr(node);
}
std::vector<Expr *> Call::expr_fields() {
  std::vector<Expr *> res;
  for (auto &x : read_args) res.push_back(&x);
  for (auto &x : write_args) res.push_back(&x);
  return res;
}
std::vector<const Expr *> Call::expr_fields() const {
  std::vector<const Expr *> res;
  for (auto &x : read_args) res.push_back(&x);
  for (auto &x : write_args) res.push_back(&x);
  return res;
}

Expr PolyFor::Make(Var iterator,
                   Expr init_val,
                   Expr condition,
                   Expr inc,
                   ForType for_type,
                   DeviceAPI device_api,
                   Expr body,
                   VectorizeInfo vectorize_info) {
  auto n        = make_shared<PolyFor>();
  n->iterator   = iterator;
  n->init       = init_val;
  n->condition  = condition;
  n->inc        = inc;
  n->device_api = device_api;
  n->body       = body;
  n->set_for_type(for_type);
  n->set_vectorize_info(vectorize_info);

  if (n->is_vectorized()) CHECK(n->vectorize_info().valid());

  return Expr(n);
}
std::vector<Expr *> PolyFor::expr_fields() { return {&init, &condition, &inc, &body}; }
std::vector<const Expr *> PolyFor::expr_fields() const { return {&init, &condition, &inc, &body}; }

Expr PolyFor::extent() const {
  auto nodes = CollectIRNodes(condition, [&](const Expr *e) {
    return e->As<NE>() ||   //
           e->As<EQ>() ||   //
           e->As<Min>() ||  //
           e->As<Max>();
  });

  if (nodes.empty()) {
    return Expr();
  }

  auto *le_n = condition.As<LE>();
  auto *lt_n = condition.As<LT>();
  if (!(le_n || lt_n)) return Expr();

  if (le_n) {
    if (le_n->a() != Expr(iterator)) return Expr();
    auto *le_b_int = le_n->b().As<IntImm>();
    if (le_b_int) return Expr(make_shared<IntImm>(Int(32), le_b_int->value + 1));
    return Add::Make(le_n->b(), Expr(1));
  }

  if (lt_n) {
    if (lt_n->a() != Expr(iterator)) return Expr();
    return lt_n->b();
  }
  return Expr();
}

bool Var::operator==(const Var &o) const { return o->name == operator->()->name; }
bool Var::operator!=(const Var &o) const { return !(*this == o); }

Var &Var::operator=(_Var_ *x) {
  *this = Var(x);
  return *this;
}

Var &Var::operator=(const _Var_ *x) {
  *this = x->Copy();
  return *this;
}

Expr Load::Make(Expr tensor, const std::vector<Expr> &indices) {
  CHECK(tensor->type().valid());
  CHECK(!indices.empty());
  for (auto &idx : indices) CHECK_EQ(idx.type().ElementOf(), Int(32));
  auto node     = make_shared<Load>();
  node->tensor  = tensor;
  node->indices = indices;
  return Expr(node);
}
Type Load::type() const {
  CHECK(tensor.defined());
  CHECK(tensor.type().valid());

  int lanes = 0;
  for (auto &idx : indices) lanes = std::max(lanes, idx.type().lanes());
  return tensor.type().ElementOf().with_lanes(lanes);
}

std::vector<Expr *> Load::expr_fields() {
  std::vector<Expr *> exprs({&tensor});
  for (auto &idx : indices) exprs.push_back(&idx);
  return exprs;
}

std::vector<const Expr *> Load::expr_fields() const {
  std::vector<const Expr *> exprs({&tensor});
  for (auto &idx : indices) exprs.push_back(&idx);
  return exprs;
}

Expr Load::index() const {
  if (is_addr_tensor()) {
    auto *tensor_n = tensor.As<_Tensor_>();
    CHECK(tensor_n);
    Expr res = common::IndiceToAbsOffset(tensor_n->shape, indices);
    optim::Simplify(&res);
    return res;
  } else {
    CHECK_EQ(indices.size(), 1UL);
    return indices[0];
  }
}

const std::string &Load::name() const {
  auto *t = tensor.As<ir::_Tensor_>();
  CHECK(t);
  return t->name;
}

bool LoadStoreAddrMnger::is_addr_tensor() const { return tensor.As<_Tensor_>(); }
bool LoadStoreAddrMnger::is_addr_scalar() const { return !is_addr_tensor(); }

Expr Ramp::Make(Expr base, Expr stride, int lanes) {
  CHECK(base.defined());
  CHECK(stride.defined());
  CHECK(base.type().valid());
  CHECK(stride.type().valid());
  CHECK_EQ(stride.type(), Int(32));
  CHECK_GT(lanes, 0);

  auto *n   = make_shared<Ramp>();
  n->base   = base;
  n->stride = stride;
  n->lanes  = lanes;
  Type type(base.type().type(), base.type().bits(), lanes);
  n->set_type(type);
  return Expr(n);
}

Expr Broadcast::Make(Expr value, int lanes) {
  CHECK(value.defined());
  CHECK(value.type().valid());

  auto *n  = make_shared<Broadcast>();
  n->value = value;
  n->lanes = lanes;

  Type type(value.type().type(), value.type().bits(), lanes);
  n->set_type(type);

  return Expr(n);
}

Type Broadcast::type() const { return value.type().ElementOf().with_lanes(lanes); }

Expr Sum::Make(const std::vector<Expr> &vs) {
  CHECK(!vs.empty());
  if (vs.size() == 1) return vs.front();

  auto *n   = make_shared<Sum>();
  auto type = vs.front().type();
  for (auto &v : vs) CHECK_EQ(v.type(), type) << vs.front() << " " << v;

  n->operands() = vs;

  n->set_type(vs.front()->type());

  return Expr(n);
}

Expr Product::Make(const std::vector<Expr> &vs) {
  CHECK_GE(vs.size(), 1);

  auto *n   = make_shared<Product>();
  auto type = vs.front().type();
  for (auto &v : vs) CHECK_EQ(v.type(), type);

  n->operands() = vs;

  n->set_type(vs.front()->type());

  return Expr(n);
}

Expr FracOp::Make(Expr n, Expr d) {
  auto *node = make_shared<FracOp>();
  node->a()  = n;
  node->b()  = d;
  return Expr(node);
}

Expr Power::Make(Expr n, Expr d) {
  auto *node          = make_shared<Power>();
  node->operands()[0] = n;
  node->operands()[1] = d;
  CHECK(d.type().is_int());

  node->set_type(n->type());

  return Expr(node);
}

lang::Module _Module_::Make(const std::string &name, Target target) {
  auto n    = make_shared<_Module_>();
  n->name   = name;
  n->target = target;
  return lang::Module(n);
}

Expr PrimitiveNode::Make(const std::string &name, const std::map<std::string, attr_t> &attrs) {
  auto *n  = make_shared<PrimitiveNode>();
  n->name  = name;
  n->attrs = attrs;
  return Expr(n);
}

}  // namespace ir
}  // namespace cinn
